package cpu

import (
	"fmt"
	"time"

	"metrics-collect/pkg/plugins"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/shirou/gopsutil/v3/cpu"
)

type CPUPlugin struct {
	*plugins.BasePlugin
	collectPerCPU bool
	cpuUsage      *prometheus.GaugeVec
	cpuTimes      *prometheus.GaugeVec
}

func NewCPUPlugin() *CPUPlugin {
	p := &CPUPlugin{
		BasePlugin: plugins.NewBasePlugin("cpu"),
	}

	// Initialize metrics
	p.cpuUsage = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "cpu_usage_percent",
			Help: "Current CPU usage percentage",
		},
		[]string{"cpu"},
	)

	p.cpuTimes = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "cpu_time_seconds",
			Help: "CPU time spent in different modes",
		},
		[]string{"cpu", "mode"},
	)

	p.AddMetric(p.cpuUsage)
	p.AddMetric(p.cpuTimes)

	return p
}

func (p *CPUPlugin) Init(config map[string]interface{}) error {
	if collectPerCPU, ok := config["collect_per_cpu"].(bool); ok {
		p.collectPerCPU = collectPerCPU
	}
	return nil
}

func (p *CPUPlugin) Start() error {
	// Start collection in a goroutine
	go func() {
		for {
			if err := p.Collect(); err != nil {
				// TODO: Use proper logging
				fmt.Printf("Error collecting CPU metrics: %v\n", err)
			}
			time.Sleep(p.GetInterval())
		}
	}()
	return nil
}

func (p *CPUPlugin) Stop() error {
	// Nothing to do for now
	return nil
}

func (p *CPUPlugin) Collect() error {
	if p.collectPerCPU {
		return p.collectPerCPUMetrics()
	}
	return p.collectTotalCPUMetrics()
}

func (p *CPUPlugin) collectPerCPUMetrics() error {
	percentages, err := cpu.Percent(0, true)
	if err != nil {
		return fmt.Errorf("failed to get CPU percentages: %v", err)
	}

	for i, percentage := range percentages {
		p.cpuUsage.WithLabelValues(fmt.Sprintf("cpu%d", i)).Set(percentage)
	}

	times, err := cpu.Times(true)
	if err != nil {
		return fmt.Errorf("failed to get CPU times: %v", err)
	}

	for i, cpuTime := range times {
		cpuLabel := fmt.Sprintf("cpu%d", i)
		p.cpuTimes.WithLabelValues(cpuLabel, "user").Set(cpuTime.User)
		p.cpuTimes.WithLabelValues(cpuLabel, "system").Set(cpuTime.System)
		p.cpuTimes.WithLabelValues(cpuLabel, "idle").Set(cpuTime.Idle)
	}

	return nil
}

func (p *CPUPlugin) collectTotalCPUMetrics() error {
	percentages, err := cpu.Percent(0, false)
	if err != nil {
		return fmt.Errorf("failed to get CPU percentages: %v", err)
	}

	if len(percentages) > 0 {
		p.cpuUsage.WithLabelValues("total").Set(percentages[0])
	}

	times, err := cpu.Times(false)
	if err != nil {
		return fmt.Errorf("failed to get CPU times: %v", err)
	}

	if len(times) > 0 {
		p.cpuTimes.WithLabelValues("total", "user").Set(times[0].User)
		p.cpuTimes.WithLabelValues("total", "system").Set(times[0].System)
		p.cpuTimes.WithLabelValues("total", "idle").Set(times[0].Idle)
	}

	return nil
}
